#include "freertos/FreeRTOS.h"
#include "esp_wifi.h"
#include "esp_wifi_types.h"
#include "esp_system.h"
#include "esp_event.h"
#include "nvs_flash.h"
#include "driver/gpio.h"
#include "lwip/ip6.h" 
#include "lwip/netif.h"

#define LED_GPIO_PIN                     5
#define WIFI_CHANNEL_SWITCH_INTERVAL  (500)
#define WIFI_CHANNEL_MAX               (13)
#define MAC_RECORD_MAX 100  // 最多记录 100 个 MAC 地址
#define DUPLICATE_TIME_THRESHOLD 60000  // 10 秒的时间阈值，单位为毫秒

typedef struct {
    char mac_str[18];
    uint32_t first_seen_time;
} MacRecord;

MacRecord mac_records[MAC_RECORD_MAX];
int mac_record_count = 0;

extern "C" int lwip_hook_ip6_input(struct pbuf *p, struct netif *inp) __attribute__((weak));
extern "C" int lwip_hook_ip6_input(struct pbuf *p, struct netif *inp) {
  if (ip6_addr_isany_val(inp->ip6_addr[0].u_addr.ip6)) {
    // We don't have an LL address -> eat this packet here, so it won't get accepted on input netif
    pbuf_free(p);
    return 1;
  }
  return 0;
}

uint8_t level = 0, channel = 1;

static wifi_country_t wifi_country = {.cc="CN", .schan = 1, .nchan = 13}; //Most recent esp32 library struct

typedef struct {
  unsigned frame_ctrl:16;
  unsigned duration_id:16;
  uint8_t addr1[6]; /* receiver address */
  uint8_t addr2[6]; /* sender address */
  uint8_t addr3[6]; /* filtering address */
  unsigned sequence_ctrl:16;
  uint8_t addr4[6]; /* optional */
} wifi_ieee80211_mac_hdr_t;

typedef struct {
  wifi_ieee80211_mac_hdr_t hdr;
  uint8_t payload[0]; /* network data ended with 4 bytes csum (CRC32) */
} wifi_ieee80211_packet_t;

// 新的事件处理函数
static void event_handler(void* arg, esp_event_base_t event_base, int32_t event_id, void* event_data)
{
    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START) {
        // 处理STA启动事件
    }
}

static void wifi_sniffer_init(void);
static void wifi_sniffer_set_channel(uint8_t channel);
static const char *wifi_sniffer_packet_type2str(wifi_promiscuous_pkt_type_t type);
static void wifi_sniffer_packet_handler(void *buff, wifi_promiscuous_pkt_type_t type);

void wifi_sniffer_init(void)
{
    nvs_flash_init();
    // 初始化网络接口
    esp_netif_init();

    // 初始化事件循环
    ESP_ERROR_CHECK(esp_event_loop_create_default());

    // 注册WiFi事件处理函数
    esp_event_handler_instance_t instance_any_id;
    esp_event_handler_instance_t instance_got_ip;
    ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &event_handler, NULL, &instance_any_id));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &event_handler, NULL, &instance_got_ip));

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK( esp_wifi_init(&cfg) );
    ESP_ERROR_CHECK( esp_wifi_set_country(&wifi_country) ); /* set country for channel range [1, 13] */
    ESP_ERROR_CHECK( esp_wifi_set_storage(WIFI_STORAGE_RAM) );
    ESP_ERROR_CHECK( esp_wifi_set_mode(WIFI_MODE_NULL) );
    ESP_ERROR_CHECK( esp_wifi_start() );
    esp_wifi_set_promiscuous(true);
    esp_wifi_set_promiscuous_rx_cb(&wifi_sniffer_packet_handler);
}

void wifi_sniffer_set_channel(uint8_t channel)
{
    esp_wifi_set_channel(channel, WIFI_SECOND_CHAN_NONE);
}

const char * wifi_sniffer_packet_type2str(wifi_promiscuous_pkt_type_t type)
{
    switch(type) {
    case WIFI_PKT_MGMT: return "MGMT";
    case WIFI_PKT_DATA: return "DATA";
    default:  
    case WIFI_PKT_MISC: return "MISC";
    }
}

static  void wifi_sniffer_packet_handler(void* buff, wifi_promiscuous_pkt_type_t type)
{
    if (type != WIFI_PKT_MGMT)
        return;

    const wifi_promiscuous_pkt_t *ppkt = (wifi_promiscuous_pkt_t *)buff;
    const wifi_ieee80211_packet_t *ipkt = (wifi_ieee80211_packet_t *)ppkt->payload;
    const wifi_ieee80211_mac_hdr_t *hdr = &ipkt->hdr;

    char mac_str[18];
    snprintf(mac_str, sizeof(mac_str), "%02x:%02x:%02x:%02x:%02x:%02x",
             hdr->addr2[0], hdr->addr2[1], hdr->addr2[2],
             hdr->addr2[3], hdr->addr2[4], hdr->addr2[5]);

    uint32_t current_time = esp_timer_get_time() / 1000; // 获取当前时间，单位为毫秒

    // 检查 MAC 地址是否已经记录，并且距离首次出现的时间是否超过 10 秒
    for (int i = 0; i < mac_record_count; i++) {
        if (strcmp(mac_records[i].mac_str, mac_str) == 0) {
            if (current_time - mac_records[i].first_seen_time < DUPLICATE_TIME_THRESHOLD) {
                return; // 10 秒内重复出现，不输出
            } else {
                mac_records[i].first_seen_time = current_time; // 更新首次出现时间
                break;
            }
        }
    }

    // 如果 MAC 地址未记录，添加到记录中
    if (mac_record_count < MAC_RECORD_MAX) {
        strcpy(mac_records[mac_record_count].mac_str, mac_str);
        mac_records[mac_record_count].first_seen_time = current_time;
        mac_record_count++;
    }

    // 尝试提取 SSID 信息
    char ssid[33] = {0};
    if ((ipkt->hdr.frame_ctrl & 0x000F) == 8) { // 信标帧
        const  uint8_t *ptr = ipkt->payload;
        while (ptr < ipkt->payload + ppkt->rx_ctrl.sig_len - 4) {
            if (*ptr == 0) { // SSID 标签
                uint8_t len = *(ptr + 1);
                if (len > 0) {
                    memcpy(ssid, ptr + 2, len);
                    ssid[len] = '\0';
                }
                break;
            }
            ptr += *(ptr + 1) + 2;
        }
    }

    printf("Found phone MAC address: %s, CHAN=%02d, RSSI=%02d, SSID=%s\n",
           mac_str, ppkt->rx_ctrl.channel, ppkt->rx_ctrl.rssi, ssid);
}

// the setup function runs once when you press reset or power the board
void setup() {
    // initialize digital pin 5 as an output.
    Serial.begin(115200);
    delay(10);
    wifi_sniffer_init();
    pinMode(LED_GPIO_PIN, OUTPUT);
}

// the loop function runs over and over again forever
void loop() {
    //Serial.print("inside loop");
    delay(1000); // wait for a second
    
    if (digitalRead(LED_GPIO_PIN) == LOW)
        digitalWrite(LED_GPIO_PIN, HIGH);
    else
        digitalWrite(LED_GPIO_PIN, LOW);
    vTaskDelay(WIFI_CHANNEL_SWITCH_INTERVAL / portTICK_PERIOD_MS);
    wifi_sniffer_set_channel(channel);
    channel = (channel % WIFI_CHANNEL_MAX) + 1;
}